Hi @acforvs, many thanks for your PR and apologies for the belated review!
It looks good to me; I left only a few comments. Could you also add tests for the rest of helpers from utils.py, so we could close the original issue?

Hi @hbq1 , thank you for the review!

Sure thing, I'll add the tests for the remaining functions, but I have some questions about them:

1. The docstring for set_diags states that the shape of new_diags should be NxD. However, new_diags is only used in one place where it is flattened.

As a result, examples like

a = jnp.ones(shape=(3, 2, 2))
diags = jnp.array([[[3, 3, 3, 3, 3, 3]]])
print(a.shape, diags.shape)  # prints (3, 2, 2) (1, 1, 6)
optax._src.utils.set_diags(a, diags)

or

a = jnp.ones(shape=(3, 2, 2))
diags = jnp.array([[3, 3, 3], [3, 3, 3]])
print(a.shape, diags.shape)  # prints (3, 2, 2) (2, 3)
optax._src.utils.set_diags(a, diags)

work just fine

Is this by design or should we raise an error in case of a shape mismatch?
Currently, the function is used internally in two places: 1, 2. The comments specify that the dimension of the passed variable is indeed NxD there, so I'd suggest adding another assert. What are your thoughts?

2. cast_tree uses .astype which wouldn't work for simple datatypes, such as

tree = dict(a=2.5, b=dict(c=-2.5))
tree = jax.tree_util.tree_map(lambda x : x, tree)
optax._src.utils.cast_tree(tree, int)

Would it be okay to only test this function for dicts of jnp.arrays of differerent dtypes? Something like

tree = dict(a=jnp.array(2.5), b=dict(c=jnp.array(-2.5)))
tree = jax.tree_util.tree_map(lambda x : x, tree)
optax._src.utils.cast_tree(tree, int)

Great, thank you!

1. It is used in some internal code, so the correct behaviour would be to raise an exception in case of shape mismatch; it'd be great if you add this to the implementation!
2. Testing only with trees of jnp.arrays would suffice.

I've fixed the errors and added tests for the rest of the utils.py functions.
I also changed assert to a ValueError in set_diags. Please let me know how it looks